### loop over thresholds for RDITS
run_no_bal <- function(threshold){
  library(data.table)
  library(tidyverse)
  library(rdrobust)
  print(threshold)
  ##keep the observations within the threshold closest to election day
  ## using data.table notation to speed this up
  set_pre <- dists[dists$dist <= threshold & dists$pre,
                   .SD[date %in% max(date)], by=list(year, GEOID)]
  set_pre <- set_pre[, .SD[1], by = .(year, GEOID)] %>% 
    mutate(treated = T)
  
  ##keep the observations within the threshold closest to election day
  set_post <- dists[dists$dist <= threshold & !dists$pre,
                    .SD[date %in% min(date)], by=list(year, GEOID)]
  set_post <- set_post[!(paste0(set_post$GEOID, set_post$year) %in% paste0(set_pre$GEOID, set_pre$year)),
                       .SD[1], by = .(year, GEOID)] %>% 
    mutate(treated = F)
  
  full_treat <- bind_rows(set_pre, set_post) %>% 
    select(GEOID, id, date, dist, year, turnout,
           median_income, nh_white, nh_black, median_age, pop_dens, latino, asian,
           some_college, turnout_pre, treated) %>% 
    mutate(d2 = ifelse(year == "2020", as.integer(date - as.Date("2020-11-03")),
                       as.integer(date - as.Date("2016-11-08"))),
           t16 = year == "2016")

  ########################################
  
  l <- rdrobust(y = full_treat$turnout, x = full_treat$d2, p = 1, c = 0, cluster = full_treat$id,
                covs = select(full_treat,
                              latino, nh_white, asian,
                              nh_black, median_income, median_age,
                              pop_dens, turnout_pre, t16, some_college))
  
  f <- tibble(coef = l$coef,
              se = l$se, 
              pv = l$pv,
              p = threshold,
              n = l$N_h[1] + l$N_h[2],
              bw = l[["bws"]][1],
              u = l[["ci"]][,2],
              l = l[["ci"]][,1])
}

cl <- makeCluster(12)  
registerDoParallel(cl)

dists <- readRDS("temp/shooting_demos.rds") |>
  mutate(date = as.Date(date),
         id = id3, ## id (and distance, in next line) of nearest 3-victim shooting
         dist = dist3) |>
  filter(!is.na(id),
         id %in% (readRDS("temp/geocoded_shootings_from_gva.rds") |>
                    filter(home + dv == 0))$id ## drop if in the home or domestic violence
  )

clusterExport(cl, list("run_no_bal", "dists"))
out <- rbindlist(parLapply(cl, c(0.25, 0.5, 0.75, seq(1, 10, 1)), fun = run_no_bal))

saveRDS(out, "temp/three_vics_out_data.rds")
out3 <- readRDS("temp/three_vics_out_data.rds")
out3$estimate <- rep(c('Traditional','Bias-Adjusted','Robust'),nrow(out3)/3)

out3 <- mutate_at(out3, vars(coef, l, u), ~. * -1)

out3$estimate <- factor(out3$estimate, levels = c('Traditional','Bias-Adjusted','Robust'))


##################
##################
##################

dists <- readRDS("temp/shooting_demos.rds") |> 
  mutate(date = as.Date(date),
         id = id5, ## id (and distance, in next line) of nearest 5-victim shooting
         dist = dist5) |> 
  filter(!is.na(id),
         id %in% (readRDS("temp/geocoded_shootings_from_gva.rds") |> 
                    filter(home + dv == 0))$id ## drop if in the home or domestic violence
  )

clusterExport(cl, list("run_no_bal", "dists"))
out <- rbindlist(parLapply(cl, c(0.25, 0.5, 0.75, seq(1, 10, 1)), fun = run_no_bal))

saveRDS(out, "temp/five_vics_out_data.rds")
out5 <- readRDS("temp/five_vics_out_data.rds")
out5$estimate <- rep(c('Traditional','Bias-Adjusted','Robust'),nrow(out5)/3)

out5 <- mutate_at(out5, vars(coef, l, u), ~. * -1)

out5$estimate <- factor(out5$estimate, levels = c('Traditional','Bias-Adjusted','Robust'))


out <- bind_rows(mutate(out3, model = "Three or More Victims"),
                 mutate(out5, model = "Five or More Victims"))

out$model <- factor(out$model, levels = c("Three or More Victims",
                                          "Five or More Victims"))

## figure s16
different_dists <- ggplot(filter(out, estimate == "Robust"),
                          aes(x = p, y = coef, ymin = l, ymax = u)) +
  facet_grid(~model) +
  geom_point() +
  geom_errorbar(width = 0) + 
  theme_bc(base_family = "LM Roman 10", base_size = 10,
           plot.margin = margin()) +
  geom_hline(yintercept = 0, linetype = "dashed") +
  scale_x_continuous(breaks = seq(0, 10, 2)) +
  labs(y = "Local average treatment effect", x = "Radius around shooting (miles)")
different_dists

fwrite(out |> 
         filter(estimate == "Robust") |> 
         select(victim_count = model,
                distance = p,
                estimate = coef,
                lower_bound = l,
                upper_bound = u), "data_for_figures/figure_A10.csv")

saveRDS(different_dists, "temp/different_dists_vic_rob.rds")

for(i in outpaths){
  ggsave(paste0(i, "vic_rob_results.png"), height = 4, width = 6, units = "in")
}

